"""
准备数据集
"""
from torch.utils.data import DataLoader, Dataset
import torch

import config


class ChatData(Dataset):
    def __init__(self):
        self.f_input = open(config.chatbot_input_by_word_path, encoding="UTF-8").readlines()
        self.f_target = open(config.chatbot_target_by_word_path, encoding="UTF-8").readlines()
        assert len(self.f_input) == len(self.f_target), "输入值和目标值数量不匹配"

    def __getitem__(self, index):
        input_data = self.f_input[index].strip().split()
        target_data = self.f_target[index].strip().split()
        input_length = len(input_data) if len(input_data) <= config.chatbot_max_len else \
            config.chatbot_max_len
        target_length = len(target_data) if len(target_data) <= config.chatbot_max_len + 1 else \
            config.chatbot_max_len + 1

        input_length = 1 if input_length <= 0 else input_length
        target_length = 1 if target_length <= 0 else target_length

        return input_data, target_data, input_length, target_length

    def __len__(self):
        return len(self.f_input)


def collate_fn(batch):
    """

    :param batch: [(input_data, target_data, input_length, target_length),
                    (input_data, target_data, input_length, target_length), ...]
    :return:
    """
    batch = sorted(batch, key=lambda x: x[2], reverse=True)
    input_data, target_data, input_length, target_length = zip(*batch)
    input_data = torch.LongTensor([config.chatbot_ws_input_by_word_model.
                                  transform(i, max_length=config.chatbot_max_len)
                                   for i in input_data])
    target_data = torch.LongTensor([config.chatbot_ws_target_by_word_model.
                                    transform(i, max_length=config.chatbot_max_len,
                                              add_eos=True) for i in target_data])

    # 注意: 这里的input_length和target_length是实际input_data和target_data的长度，不受chatbot_max_len的影响
    input_length = torch.LongTensor(input_length)
    target_length = torch.LongTensor(target_length)

    return input_data, target_data, input_length, target_length


train_data_loader = DataLoader(ChatData(), batch_size=config.chatbot_batch_size, shuffle=True,
                               collate_fn=collate_fn)




